import math

from torch.utils.data import DataLoader

import torch.nn.functional as F

from model import *

"""
训练
"""

"""
reshape(-1) 的含义如下：
功能：将张量重新塑形为一维数组
参数解释：
-1 表示自动推断该维度的大小
PyTorch会根据张量中元素的总数自动计算出合适的维度
"""


def train_one_epoch(model: Seq2seqTransformer, loader: DataLoader, optimizer: torch.optim.Optimizer):
    model.train()
    total, correct, tok = 0.0, 0.0, 0
    for src, tgt in loader:
        src, tgt = src.to(DEVICE), tgt.to(DEVICE)

        tgt_inp = tgt[:, :-1]
        tgt_out = tgt[:, 1:]

        src_pad = (src == PAD_IDX)  # [B,S] True=pad
        tgt_pad = (tgt_inp == PAD_IDX)  # [B,T-1]
        tgt_mask = generate_square_subsequent_mask(tgt_inp.size(1)).to(DEVICE)  # [T-1,T-1]

        logits = model(src, tgt_inp, src_pad, tgt_pad, tgt_mask)  # [B,T-1,V]
        # 计算模型输出 logits 与目标序列 tgt_out 之间的交叉熵损失。
        # 将 logits 和目标展平为一维，忽略填充值（PAD_IDX）的影响，得到平均损失值用于反向传播。
        loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), tgt_out.reshape(-1), ignore_index=PAD_IDX)

        optimizer.zero_grad()
        loss.backward()
        # 用于梯度裁剪，防止在反向传播过程中梯度过大导致训练不稳定。
        # 它将模型参数的梯度范数限制在1.0以内，确保梯度更新更加平稳，有助于提升模型训练的稳定性和收敛速度。
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        """
        在不计算梯度的情况下评估模型性能：
        pred = logits.argmax(-1)：获取模型预测的标签；
        mask = (tgt_out != PAD_IDX)：生成非填充位置的掩码；
        correct += ((pred == tgt_out) & mask).sum().item()：统计非填充位置中预测正确的 token 数量；
        tok += mask.sum().item()：累计非填充 token 总数；
        total += loss.item() * mask.sum().item()：加权累计损失，用于后续计算困惑度（PPL）。
        """
        with torch.no_grad():
            pred = logits.argmax(-1)  # [B,T-1]
            mask = (tgt_out != PAD_IDX)
            correct += ((pred == tgt_out) & mask).sum().item()
            tok += mask.sum().item()
            total += loss.item() * mask.sum().item()

    ppl = math.exp(total / max(tok, 1))
    acc = correct / max(tok, 1)
    return ppl, acc


"""
推理
"""


@torch.no_grad()
def evaluate(model, loader):
    model.eval()
    total, correct, tok = 0.0, 0.0, 0
    for src, tgt in loader:
        src, tgt = src.to(DEVICE), tgt.to(DEVICE)
        tgt_inp = tgt[:, :-1]
        tgt_out = tgt[:, 1:]
        src_pad = (src == PAD_IDX)
        tgt_pad = (tgt_inp == PAD_IDX)
        tgt_mask = generate_square_subsequent_mask(tgt_inp.size(1)).to(DEVICE)
        logits = model(src, tgt_inp, src_pad, tgt_pad, tgt_mask)
        loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), tgt_out.reshape(-1), ignore_index=PAD_IDX)
        pred = logits.argmax(-1)
        mask = (tgt_out != PAD_IDX)
        correct += ((pred == tgt_out) & mask).sum().item()
        tok += mask.sum().item()
        total += loss.item() * mask.sum().item()

    ppl = math.exp(total / max(tok, 1))
    acc = correct / max(tok, 1)
    return ppl, acc
